import gradio as gr
import re
import paddle
import shutil
from paddlenlp.transformers import BertTokenizer
from PoetryGen import PoetryGen
from Poetry_Bert_Model import PoetryBertModel
from gradio_client import Client


client = Client("https://s5k.cn/api/v1/studio/gally16/Bert-VITS21.x/gradio/")
bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
net = PoetryBertModel('bert-base-chinese', 128)
model = paddle.Model(net)
model.load('./checkpoint_10/final')
poetry_gen = PoetryGen(model, bert_tokenizer)

# 遇到标点符号换行的工具函数
def poetry_show(poetry):
    result = ""
    pattern = r"([，。；？])"
    text = re.sub(pattern, r'\1 ', poetry)
    for p in text.split():
        if p:
            # print(p)
            result += p
            result += "\n"
    return result[:-1]

# 拿到wav文件路径
'''
Gradio 中，默认情况下，临时文件会存储在 /tmp 目录中。然而，有时我们的应用程序可能无法访问 /tmp 目录，这可能会导致文件操作失败或出现其他异常情况。
'''
def cut_wav(source_folder):

    gradio_id = source_folder.split('\\')[-2]
    # print("gradio_id",gradio_id)
    # 目标文件夹路径
    destination_folder = r'./gradio_wav'
    # 移动文件夹
    # print("source_folder[:-9],",source_folder[:-9],)
    shutil.move(source_folder[:-9], destination_folder)

    wav_path = destination_folder + "/" +gradio_id + "/audio.wav"
    return wav_path

def gen_poem(x, style, mode, isread=False, role="八重神子"):
    global mode_12
    if mode == "藏头诗" and x is not None:
        mode_12 = 2
    elif mode == "续写":
        mode_12 = 1
    poetry = poetry_gen.generate(head=x, mode=mode_12, style=style)
    if isread:
        response = client.predict(
            poetry,  # str  in '请输入简体中文文案' Textbox component
            role,
            # Literal[钟离, 凝光, 克列门特, 海芭夏, 龙二, 埃尔欣根, 松浦, 安柏, 浣溪, 霄翰, 克拉拉, 萨齐因, 艾伯特, 九条镰治, 莫娜, 奥列格, 早柚, 芭芭拉, 达达利亚, 柯莱, 阿兰, 纳比尔, 空, 恕筠, 佩拉, 阿晃, 诺艾尔, 「博士」, 三月七, 布洛妮娅, 宛烟, 佐西摩斯, 七七, 托克, 班尼特, 阿尔卡米, 阿洛瓦, 伊迪娅, 羽生田千鹤, 阿祇, 舒伯特, 言笑, 埃舍尔, 明曦, 埃德, 温迪, 桑博, 景元, 云堇, 砂糖, 知易, 米卡, 宵宫, 艾莉丝, 罗刹, 嘉玛, 绿芙蓉, 掇星攫辰天君, 「信使」, 青镞, 西拉杰, 可可利亚, 巴达维, 杜拉夫, 帕斯卡, 辛焱, 芙宁娜, 阿圆, 荒泷一斗, 希露瓦, 雷电将军, 柊千里, 旁白, 素裳, 埃泽, 重云, 上杉, 枫原万叶, 史瓦罗, 夏洛蒂, 鹿野奈奈, 优菈, 帕姆, 青雀, 白术, 式大将, 螺丝咕姆, 烟绯, 阿扎尔, 坎蒂丝, 费斯曼, 大慈树王, 长生, 流浪者, 笼钓瓶一心, 凯亚, 琴, 阿娜耶, 莎拉, 哲平, 绮良良, 海妮耶, 行秋, 托马, 克罗索, 查尔斯, 毗伽尔, 赛诺, 德沃沙克, 开拓者(女), 博来, 丹枢, 天叔, 甘雨, 纳西妲, 陆行岩本真蕈·元素生命, 菲谢尔, 阿守, 公输师傅, 蒂玛乌斯, 刻晴, 女士, 莺儿, 艾尔海森, 迪奥娜, 久利须, 卡芙卡, 妮露, 莱依拉, 神里绫华, 伊利亚斯, 艾丝妲, 多莉, 「大肉丸」, 迈蒙, 慧心, 卢卡, 回声海螺, anzai, 「散兵」, 奥兹, 金人会长, 老孟, 罗莎莉亚, 「女士」, 雷泽, 萨赫哈蒂, 埃洛伊, 塞塔蕾, 银狼, 萍姥姥, 符玄, 埃勒曼, 元太, 白露, 珊瑚宫心海, 彦卿, 霍夫曼, 虎克, 八重神子, 姬子, 爱贝尔, 斯坦利, 玲可, 深渊使徒, 提纳里, 香菱, 神里绫人, 杰帕德, 瓦尔特, 玛格丽特, 莫塞伊思, 「白老先生」, 开拓者(男), 迪希雅, 艾文, 驭空, 卡维, 影, 娜塔莎, 可莉, 昆钧, 阿贝多, 岩明, 菲米尼, 拉齐, 娜维娅, 伦纳德, 鹿野院平藏, 欧菲妮, 荧, 石头, 纯水精灵？, 迪娜泽黛, 申鹤, 浮游水蕈兽·元素生命, 黑塔, 天目十五, 玛乔丽, 晴霓, 魈, 瑶瑶, 珊瑚, 丹恒, 久岐忍, 丹吉尔, 塞琉斯, 迪卢克, 凯瑟琳, 留云借风真君, 恶龙, 玛塞勒, 悦, 希儿, 吴船长, 沙扎曼, 北斗, 博易, 斯科特, 百闻, 丽莎, 安西, 迈勒斯, 镜流, 琳妮特, 戴因斯雷布, 林尼, 康纳, 夜兰, 阿拉夫, 那维莱特, 「公子」, 卡波特, 阿佩普, 半夏, 塔杰·拉德卡尼, 胡桃, 大毫, 深渊法师, 五郎, 九条裟罗, 拉赫曼, 田铁嘴, 嘉良, 阿巴图伊, 常九爷, 刃, 爱德琳, 派蒙, 珐露珊, 停云]  in '角色' Dropdown component
            0.2,  # float (numeric value between 0 and 1) in '语调调节' Slider component
            0.6,  # float (numeric value between 0.1 and 2) in '感情调节' Slider component
            0.8,  # float (numeric value between 0.1 and 2) in '音素长度' Slider component
            1.2,  # float (numeric value between 0.1 and 2) in '生成时长' Slider component
            api_name="/predict_1"
        )
        # 使用示例
        input_wav_path = str(response)
        print("input_wav_path:",input_wav_path)
        output_wav_path = cut_wav(input_wav_path)
        # print("output_wav_path:",output_wav_path)
        return poetry_show(poetry), output_wav_path
    else:
        return poetry_show(poetry)


with gr.Blocks() as demo:
    gr.Markdown(''''
                # 中华诗词生成工具
                ## 作者：谢凯军
                ### 2024/7/8''')
    # 设置tab选项卡
    with gr.Tab("诗词生成"):
        with gr.Column():
            mode_choise = gr.Radio(["藏头诗", "续写"], value="续写", label="模式选择", info="想要藏头诗or续写")
            text_input1 = gr.Textbox(label="内容", lines=2)
            text_input2 = gr.Textbox(label='style', lines=1)
            text_output = gr.Textbox(label="生成的诗词")
            text_button = gr.Button("生成")

    with gr.Tab("诗词朗诵生成"):
        with gr.Row():
            mode_choise2 = gr.Radio(["藏头诗", "续写"], value="续写", label="模式选择", info="想要藏头诗or续写")
            text_input12 = gr.Textbox(label="内容", lines=6)
            is_read = gr.State(value=True)
            text_input22 = gr.Textbox(label='style', lines=6)
            text_output2 = gr.Textbox(label="生成的诗词", lines=6)
        audio_output2 = gr.Audio(label="朗诵语音",type="filepath",autoplay=True)
        text_button2 = gr.Button("生成")

    with gr.Accordion("打开查看更多", open=False):
        gr.Markdown("打开也没用,我还没写")

    text_button.click(gen_poem,
                      inputs=[text_input1, text_input2, mode_choise],
                      outputs=text_output)
    text_button2.click(gen_poem,
                       inputs=[text_input12, text_input22, mode_choise2, is_read],
                       outputs=[text_output2, audio_output2])

demo.launch()
